package com.luo.ai.assistant.etl.impl;

import com.luo.ai.assistant.etl.MyDocumentReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.Resource;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 文档读取器实现
 *
 * @author luohq
 * @date 2025-01-01 15:24
 */
public class MyDocumentReaderImpl implements MyDocumentReader {
    private static final Logger log = LoggerFactory.getLogger(MyDocumentReaderImpl.class);

    private final Resource resource;

    private final Map<String, Object> customMetadata = new HashMap<>();

    public MyDocumentReaderImpl(String resourceUrl) {
        this(new DefaultResourceLoader().getResource(resourceUrl));
    }

    public MyDocumentReaderImpl(Resource resource) {
        Objects.requireNonNull(resource, "The Spring Resource must not be null");
        this.resource = resource;
    }

    @Override
    public Map<String, Object> getCustomMetadata() {
        return this.customMetadata;
    }

    /**
     * 根据文档类型自动选择读取器
     *
     * @return List of Documents.
     */
    @Override
    public List<Document> get() {
        try {
            String filename = this.resource.getFilename().toLowerCase();
            String fileType = filename.substring(filename.lastIndexOf(".") + 1);
            switch (fileType) {
                case "json":
                    return this.readJson();
                case "md":
                    return this.readMarkdown();
                case "pdf":
                    return this.readPdfAutoDetect();
                case "doc":
                case "docx":
                case "ppt":
                case "pptx":
                case "xls":
                case "xlsx":
                case "htm":
                case "html":
                    return this.readOffice();
                default:
                    return this.readText();
            }
        } catch (Throwable e) {
            log.warn("read document exception!", e);
            return readText();
        }
    }

    @Override
    public List<Document> readJson(String... jsonKeysToUse) {
        JsonReader jsonReader = new JsonReader(this.resource, jsonKeysToUse);
        return jsonReader.get();
    }

    @Override
    public List<Document> readText() {
        TextReader textReader = new TextReader(this.resource);
        textReader.getCustomMetadata().putAll(this.customMetadata);
        return textReader.read();
    }

    @Override
    public List<Document> readMarkdown() {
        MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                .withHorizontalRuleCreateDocument(true)
                .withIncludeCodeBlock(false)
                .withIncludeBlockquote(false)
                .withAdditionalMetadata(this.customMetadata)
                .build();

        MarkdownDocumentReader reader = new MarkdownDocumentReader(this.resource, config);
        return reader.get();
    }

    @Override
    public List<Document> readPdfAutoDetect() {
        try {
            return this.readPdfByCatalog();
        } catch (Throwable ex) {
            log.info("Not support readPdfByCatalog and Change to readPdfPages - exception: {}", ex.getMessage());
            return this.readPdfPages();
        }
    }

    @Override
    public List<Document> readPdfPages() {
        //The PagePdfDocumentReader uses Apache PdfBox library to parse PDF documents
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(this.resource,
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());

        return pdfReader.read();
    }

    @Override
    public List<Document> readPdfByCatalog() {
        //The ParagraphPdfDocumentReader uses the PDF catalog (e.g. TOC) information to split the input PDF into text paragraphs
        //and output a single Document per paragraph.
        //NOTE: Not all PDF documents contain the PDF catalog.
        ParagraphPdfDocumentReader pdfReader = new ParagraphPdfDocumentReader(this.resource,
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());

        return pdfReader.read();
    }

    @Override
    public List<Document> readOffice() {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(this.resource);
        return tikaDocumentReader.read();
    }
}
